# coding: utf-8
import PIL.Image
# @Time: 2025-05-26

from utils import *
import os
import numpy as np
import commpy
from commpy.modulation import QAMModem
from commpy.channelcoding.ldpc import get_ldpc_code_params, ldpc_bp_decode, triang_ldpc_systematic_encode
from pyldpc import make_ldpc, encode, decode, get_message, ldpc_images, encode_no_noise, decode_no_noise
from pyldpc.utils_img import gray2bin, rgb2bin
import math
from tqdm import tqdm
from PIL import Image
Image.MAX_IMAGE_PIXELS = None
# multi-thread
from concurrent.futures import ThreadPoolExecutor, as_completed
# multi-pool
from multiprocessing import Pool, cpu_count
from typing import List
from numba import jit
import torch
import warnings


#+++++++++++++++++++++++++++++++++++++++++++++++++++ 信道编码 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++#
class LDPC:
    """
    LDPC 编码
    """
    def __init__(self,
                 n: int = 720,
                 d_v: int = 2,
                 rate: float = 1/2,  # LDPC1/2
                 snr: float = 5,
                 seed: int = 42
                 ):
        # codeword length
        self.n = n
        self.d_v = d_v
        self.d_c = d_v * 2 if rate == 1/2 else int(d_v / (1 - rate))
        self.snr = snr
        self.seed = seed
        self.img_bin_shape = None
        # H, G: decoding and coding matrix. H is a regular parity-check matrix with d_v ones per row
        # and d_c ones per column.
        self.H, self.G = make_ldpc(n, d_v, self.d_c, seed=seed, systematic=True, sparse=True)
        self.k = self.G.shape[1]    # message length
        pass

    def encode(self, msg):
        codeword = encode(self.G, msg, self.snr, self.seed)
        return codeword

    def encode_no_noise(self, msg):
        codeword = encode_no_noise(self.G, msg, self.seed)
        return codeword

    def decode(self, msg):
        decoded = decode(self.H, msg, self.snr, maxiter=50)
        msg_hat = get_message(self.G, decoded)
        return msg_hat

    def decode_no_noise(self, msg):
        decoded = decode_no_noise(self.H, msg, maxiter=50)
        msg_hat = get_message(self.G, decoded)
        return msg_hat


class LDPC_v2:
    def __init__(self,
                 ldpc_design_file: str = "./1440.720.txt"
                 ):
        self.param = get_ldpc_code_params(ldpc_design_file, compute_matrix=True)
        pass

    def encode(self, msg_bits):
        self.msg_len = len(msg_bits)
        return triang_ldpc_systematic_encode(msg_bits, self.param)

    def decode(self, msg_bits):
        # LDPC解码时需要将1变为-1，将0变为1
        # msg_bits[msg_bits == 1] = -1
        # msg_bits[msg_bits == 0] = 1
        ldpc_decoded_bits = ldpc_bp_decode(msg_bits.reshape(-1, order='F').astype(float), self.param, 'MSA', 10)[0][
                            :720].reshape(-1, order='F')[:self.msg_len]
        return ldpc_decoded_bits


#+++++++++++++++++++++++++++++++++++++++++++++++++++ 信号调制 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++#
class QAM:
    def __init__(self,
                 M: int = 16,   # 16-QAM
                 ):
        self.M = M
        self.qam = QAMModem(M)
        pass

    def modulate(self, bits):
        # QAM 调制
        return self.qam.modulate(bits)

    def demodulate(self, rx_symbols):
        # QAM 解调， 硬编码
        return self.qam.demodulate(rx_symbols, "hard")


#+++++++++++++++++++++++++++++++++++++++++++++++++++ 信道噪声 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++#
class Channel:
    def __init__(self,
                 channel_type: str = "awgn",
                 snr: float = 5
                 ):
        self.channel_type = channel_type
        self.snr = snr  # 单位 dB
        print("SNR: %d dB." % snr )
        pass

    def pass_channel(self, modulated_bits):
        """
        awgn: 白性加高斯噪声
        rayleigh: 瑞利信道
        rician: 莱斯信道
        """
        input_shape = modulated_bits.shape
        modulated_bits = modulated_bits.reshape(-1)

        if self.channel_type == 'awgn':
            bits_with_noise = commpy.awgn(modulated_bits, self.snr)

        elif self.channel_type == 'rayleigh':
            N = len(modulated_bits)
            h = (np.random.randn(N) + 1j * np.random.randn(N)) / np.sqrt(2)
            modulated_bits_times_channel_coefficients = modulated_bits * h
            modulated_bits_times_channel_coefficients_with_noise = commpy.awgn(modulated_bits_times_channel_coefficients,
                                                                               self.snr)
            bits_with_noise = modulated_bits_times_channel_coefficients_with_noise / h

        elif self.channel_type == 'rician':
            K = 100
            N = len(modulated_bits)
            m = np.random.randn(N, 1)
            t = np.random.randn(N, 1)
            Complex_Mat = 1j * m[1, :]
            Complex_Mat += t[:, 1]
            h = math.sqrt(K / K + 1) + math.sqrt(1 / K + 1) * Complex_Mat / math.sqrt(2)
            s = modulated_bits * h
            r = commpy.awgn(s, self.snr)
            bits_with_noise = r / h

        else:
            raise ValueError("Unsupported channel type.")

        bits_with_noise = bits_with_noise.reshape(input_shape)

        return bits_with_noise


#++++++++++++++++++++++++++++++++++++++++ 信道编码 + 信号调制 + 信道噪声 +++++++++++++++++++++++++++++++++++++++++++++++++#
class SimulateCommunication:
    def __init__(self,
                 config):
        self.config = config
        self.codec = LDPC(
            self.config.n,
            self.config.d_v,
            self.config.rate,
            self.config.snr,
            self.config.seed
        )
        # self.codec = LDPC_v2()
        self.modulate = QAM(
            self.config.M
        )
        self.channel = Channel(
            self.config.channel_type, self.config.snr
        )
        pass

    def __call__(self, bits):
        pad_len = (-len(bits)) % self.codec.k
        bits_padded = np.hstack([bits, np.random.randint(0, 2, pad_len)])
        # decoded_bits = []
        data_list = [bits_padded[i:i+self.codec.k] for i in range(0, len(bits_padded), self.codec.k)]
        time_node1 = time.time()
        # msg = bits_padded[i:i+self.codec.k]
        # msg_hat = self.process_one(msg)
        # decoded_bits.extend(msg_hat)
        # decoded_bits = self.threaded_process(data_list, max_workers=16)      # np.ndarray
        decoded_bits = self.multiprocess_ordered(data_list)     # list
        # decoded_bits = self.multiprocess_thread_ordered(data_list)
        print(f"epoch: {time.time()-time_node1} secs.")
        return decoded_bits[:len(bits)]

    def process_one(self, msg_bits):
        import time
        start = time.time()
        codeword = self.codec.encode_no_noise(msg_bits)  # code
        symbols = self.modulate.modulate(codeword)  # modulation
        rx = self.channel.pass_channel(symbols)  # pass channel
        llr = self.modulate.demodulate(rx)
        msg_hat = self.codec.decode_no_noise(llr)
        # print(f"one process: {time.time()-start} secs.")
        return msg_hat

    def threaded_process(self, data_list, max_workers=16):
        # 多线程
        results = [None] * len(data_list)
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            future_to_index = {
                executor.submit(self.process_one, data): idx
                for idx, data in enumerate(data_list)
            }

            for future_idx in as_completed(future_to_index):
                idx = future_to_index[future_idx]
                results[idx] = future_idx.result()
        return np.array(results).flatten()

    def multiprocess_ordered(self, data, num_process: int = None):
        # 多进程
        if num_process is None:
            num_process = cpu_count()
        print(f"num_process: {num_process}.")
        with Pool(processes=num_process) as pool:
            results = pool.map(self.process_one, data)
        return results

    def multiprocess_thread_ordered(self, data_list: List[List],
                                    num_process: int = None):
        # 多进程 + 多线程
        if num_process is None:
            num_process = cpu_count()
        print(f"num_process: {num_process}")
        item_process = len(data_list) // num_process
        # [[[], [], ...], [[], [], ...]]
        data_list_list = [data_list[i:i+item_process] for i in range(0, len(data_list), item_process)]
        with Pool(processes=num_process) as pool:
            results = pool.map(self.threaded_process, data_list_list)
        return results

    def start(self, bits):
        time_node1 = time.time()
        codeword = self.codec.encode(bits)
        time_node2 = time.time()
        print(f"ldpc encode: {time_node2 - time_node1} secs.")
        codeword_shape = codeword.shape
        symbols = self.modulate.modulate(codeword.reshape(-1))
        time_node3 = time.time()
        print(f"qam modulate: {time_node3-time_node2} secs.")
        rx = self.channel.pass_channel(symbols)
        time_node4 = time.time()
        print(f"pass channel: {time_node4 - time_node3} secs.")
        llr = self.modulate.demodulate(rx)
        time_node5 = time.time()
        print(f"qam demodulate: {time_node5-time_node4} secs.")
        llr = llr.reshape(codeword_shape[0], -1)
        msg_hat = self.codec.decode(llr)
        time_node6 = time.time()
        print(f"ldpc decode: {time_node6 - time_node5} secs.")
        return msg_hat


def ldpc_qam_awgn(input_signal, ldpc_param, snr=2, qam_order=16):
    import time
    binary_arr = input_signal
    """
    LDPC信道编码
    """
    # 赋给message_bits作为信道编码的输入
    message_bits = binary_arr
    # 指定LDPC编码文件
    # ldpc_design_file = './1440.720.txt'
    # 从编码文件中获取编码参数
    # param = get_ldpc_code_params(ldpc_design_file, compute_matrix=True)
    param = ldpc_param
    # 调用编码函数生成LDPC编码后的消息比特, 输入应为一维数组，输出为二维数组
    # time_node1 = time.time()
    ldpc_encoded_bits = triang_ldpc_systematic_encode(message_bits, param)
    # time_node2 = time.time()
    # print(f"LDPC encode: {time_node2 - time_node1} secs.")
    """
    QAM调制
    """
    # 记录LDPC编码后的二维数组的尺寸便于后面恢复
    first_dimension_length = ldpc_encoded_bits.shape[0]
    # 将二维转为一维
    bits = ldpc_encoded_bits.reshape(-1)
    # 实例化4QAM模型
    qam_model = commpy.QAMModem(qam_order)
    # 利用实例化的QAM模型对比特流进行调制，输入需要为一维数组, 耗时
    modulated_bits = qam_model.modulate(bits)
    # time_node3 = time.time()
    # print(f"QAM modulate: {time_node3 - time_node2} secs.")
    """+
    AWGN信道
    """
    # 对调制结果通过AWGN信道（指定SNR）
    bits_with_noise = commpy.awgn(modulated_bits, snr)
    # time_node4 = time.time()
    # print(f"AWGN: {time_node4 - time_node3} secs.")
    """
    QAM解调
    """
    # QAM解调， 耗时
    demodulated_bits = qam_model.demodulate(bits_with_noise, 'hard')
    print(np.array_equal(bits, demodulated_bits))
    # time_node5 = time.time()
    # print(f"QAM demodulate: {time_node5 - time_node4} secs.")
    """
    LDPC解码
    """
    # 将一维重新恢复为二维
    ldpc_encoded_bits = demodulated_bits.reshape(first_dimension_length, -1)
    # LDPC解码时需要将1变为-1，将0变为1
    # ldpc_encoded_bits[ldpc_encoded_bits == 1] = -1
    # ldpc_encoded_bits[ldpc_encoded_bits == 0] = 1
    ldpc_encoded_bits = 1 - 2 * ldpc_encoded_bits
    # 将上面的语句合并， 耗时
    ldpc_decoded_bits = ldpc_bp_decode(ldpc_encoded_bits.reshape(-1, order='F').astype(float), param, 'MSA', 10)[0][:720].reshape(-1, order='F')[:len(message_bits)]
    print(np.array_equal(ldpc_decoded_bits, message_bits))
    # time_node6 = time.time()
    # print(f"LDPC decode: {time_node6 - time_node5} secs.")
    return ldpc_decoded_bits


"""
LDPC信道编码 -> QAM调制 -> rayleigh信道 -> QAM解调 -> LDPC信道解码
"""
def ldpc_qam_rayleigh(input_signal, snr=2, qam_order=16, ldpc_design_file='./1440.720.txt'):

    binary_arr = input_signal
    """
    LDPC信道编码
    """
    # 赋给message_bits作为信道编码的输入
    message_bits = binary_arr
    # 指定LDPC编码文件
    # ldpc_design_file = './1440.720.txt'
    # 从编码文件中获取编码参数
    param = get_ldpc_code_params(ldpc_design_file)
    # 调用编码函数生成LDPC编码后的消息比特, 输入应为一维数组，输出为二维数组
    ldpc_encoded_bits = triang_ldpc_systematic_encode(message_bits, param)
    """
    QAM调制
    """
    # 记录LDPC编码后的二维数组的尺寸便于后面恢复
    first_dimension_length = ldpc_encoded_bits.shape[0]
    # 将二维转为一维
    bits = ldpc_encoded_bits.reshape(-1)
    # 实例化4QAM模型
    qam_model = commpy.QAMModem(qam_order)
    # 利用实例化的QAM模型对比特流进行调制，输入需要为一维数组
    modulated_bits = qam_model.modulate(bits)
    """
   瑞利信道:复高斯法实现
    """
    N = len(modulated_bits)
    # Calculate the channel coefficients
    h = (np.random.randn(N) + 1j * np.random.randn(N)) / np.sqrt(2)
    modulated_bits_times_channel_coefficients = modulated_bits * h
    modulated_bits_times_channel_coefficients_with_noise = commpy.awgn(modulated_bits_times_channel_coefficients, snr)
    # 接收端需要重新除以信道参数来估计输出
    bits_with_noise = modulated_bits_times_channel_coefficients_with_noise / h
    """
    QAM解调
    """
    # QAM解调
    demodulated_bits = qam_model.demodulate(bits_with_noise, 'hard')
    print(np.array_equal(bits, demodulated_bits))
    """
    LDPC解码
    """
    # 将一维重新恢复为二维
    ldpc_encoded_bits = demodulated_bits.reshape(first_dimension_length, -1)
    # LDPC解码时需要将1变为-1，将0变为1
    # ldpc_encoded_bits[ldpc_encoded_bits == 1] = -1
    # ldpc_encoded_bits[ldpc_encoded_bits == 0] = 1
    ldpc_encoded_bits = 1 - 2 * ldpc_encoded_bits
    # 将上面的语句合并
    ldpc_decoded_bits = ldpc_bp_decode(ldpc_encoded_bits.reshape(-1, order='F').astype(float), param, 'MSA', 10)[0][:720].reshape(-1, order='F')[:len(message_bits)]
    print(np.array_equal(ldpc_decoded_bits, message_bits))
    return ldpc_decoded_bits


def ldpc_qam_rician(input_signal, snr=2, qam_order=16, ldpc_design_file='./1440.720.txt'):
    """
    LDPC信道编码 -> QAM调制 -> rician信道 -> QAM解调 -> LDPC信道解码
    """
    K = 100
    binary_arr = input_signal
    """
    LDPC信道编码
    """
    # 赋给message_bits作为信道编码的输入
    message_bits = binary_arr
    # 指定LDPC编码文件
    # ldpc_design_file = './1440.720.txt'
    # 从编码文件中获取编码参数
    param = get_ldpc_code_params(ldpc_design_file)
    # 调用编码函数生成LDPC编码后的消息比特, 输入应为一维数组，输出为二维数组
    ldpc_encoded_bits = triang_ldpc_systematic_encode(message_bits, param)
    """
    QAM调制
    """
    # 记录LDPC编码后的二维数组的尺寸便于后面恢复
    first_dimension_length = ldpc_encoded_bits.shape[0]
    # 将二维转为一维
    bits = ldpc_encoded_bits.reshape(-1)
    # 实例化4QAM模型
    qam_model = commpy.QAMModem(qam_order)
    # 利用实例化的QAM模型对比特流进行调制，输入需要为一维数组
    modulated_bits = qam_model.modulate(bits)
    """
    莱斯信道：莱斯因子K可以在136行进行设置
    """
    # 对调制结果通过莱斯信道（指定SNR）
    N = len(modulated_bits)
    m = np.random.randn(N, 1)
    t = np.random.randn(N, 1)
    Complex_Mat = 1j * m[1, :]
    Complex_Mat += t[:, 1]
    h=math.sqrt(K/K+1)+math.sqrt(1/K+1)*Complex_Mat/math.sqrt(2)
    s=modulated_bits*h
    r = commpy.awgn(s, snr)
    bits_with_noise=r/h
    """
    QAM解调
    """
    # QAM解调
    demodulated_bits = qam_model.demodulate(bits_with_noise, 'hard')
    print(np.array_equal(bits, demodulated_bits))
    """
    LDPC解码
    """
    # 将一维重新恢复为二维
    ldpc_encoded_bits = demodulated_bits.reshape(first_dimension_length, -1)
    # LDPC解码时需要将1变为-1，将0变为1
    # ldpc_encoded_bits[ldpc_encoded_bits == 1] = -1
    # ldpc_encoded_bits[ldpc_encoded_bits == 0] = 1
    ldpc_encoded_bits = 1 - 2 * ldpc_encoded_bits
    # 将上面的语句合并
    ldpc_decoded_bits = ldpc_bp_decode(ldpc_encoded_bits.reshape(-1, order='F').astype(float), param, 'MSA', 10)[0][:720].reshape(-1, order='F')[:len(message_bits)]
    print(np.array_equal(ldpc_decoded_bits, message_bits))
    return ldpc_decoded_bits


#+++++++++++++++++++++++++++++++++++++++++++++++++++ 评估函数 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++#
def compute_ber(bits, rx_bits):
    return np.mean(bits != rx_bits)



#++++++++++++++++++++++++ LDPC -> QAM -> channel (awgn, rayleigh, rician) ++++++++++++++++++++++++++++++++++++++++++++#
# 新的通信方式
class LdpcQamChannel:
    def __init__(self,
                 save_folder,
                 snr: float = 5,
                 ldpc_rate: float = 1/2,
                 qam_m: int = 16,
                 ldpc_design_file: str = "./1440.720.txt",
                 channel_type: str = "awgn"
                 ):
        self.save_folder = save_folder  # 保存经过信道后解码的数据
        os.makedirs(save_folder, exist_ok=True)
        self.ldpc_rate = ldpc_rate
        self.qam_m = qam_m
        self.channel_type = channel_type
        self.snr = snr
        self.channel = Channel(channel_type, snr)       # 初始化信道
        self.ldpc_param = get_ldpc_code_params(ldpc_design_file)
        self.qam_model = commpy.QAMModem(self.qam_m)
        pass
    
    def transmite(self, input_signal):
        # input_signal 为二进制比特流, 一维
        message_bits = input_signal
        # LDPC 编码
        ldpc_encoded_bits = triang_ldpc_systematic_encode(message_bits, self.ldpc_param) 
        first_dimention_length = ldpc_encoded_bits.shape[0]
        bits = ldpc_encoded_bits.reshape(-1)
        # QAM 调制
        modulated_bits = self.qam_model.modulate(bits)
        # 通过信道
        bits_with_noise = self.channel.pass_channel(modulated_bits)
        # QAM 解调
        demodulated_bits = self.qam_model.demodulate(bits_with_noise, "hard")
        # print("调制与解调后的信号是否没有失真：", np.array_equal(bits, demodulated_bits))
        # LDPC 解码
        # ldpc_encoded_bits = demodulated_bits.reshape(first_dimention_length, -1)
        ldpc_tobe_decode_bits = demodulated_bits.reshape(first_dimention_length, -1)
        # LDPC 解码时需要将 1 变为 -1， 将 0 变为 1
        # ldpc_encoded_bits[ldpc_encoded_bits == 1] = -1
        #
        ldpc_tobe_decode_bits[ldpc_tobe_decode_bits == 1] = -1
        ldpc_tobe_decode_bits[ldpc_tobe_decode_bits == 0] = 1
        # ldpc_decoded_bits = ldpc_bp_decode(ldpc_encoded_bits.reshape(-1, order='F').astype(float), param, 'MSA', 10)[0][:720].reshape(-1, order='F')[:len(message_bits)]
        ldpc_decoded_bits = ldpc_bp_decode(ldpc_tobe_decode_bits.reshape(-1, order='F').astype(float), self.ldpc_param, "MSA", 10)[0][:720].reshape(-1, order='F')[:len(message_bits)]
        return ldpc_decoded_bits 

    def save_image(self, image: Image or np.ndarray,
                   name: str,
                   suffix: str = ".png"):
        if isinstance(image, PIL.Image.Image):
            image.save(os.path.join(self.save_folder, name + suffix))
        elif isinstance(image, np.ndarray):
            image.tofile(os.path.join(self.save_folder, name + suffix))


def save_image(image: Image or np.ndarray,
               save_folder: str,
               name: str,
               suffix: str = ".png"):
    if isinstance(image, PIL.Image.Image):
        image.save(os.path.join(save_folder, name + suffix))
    elif isinstance(image, np.ndarray):
        image.tofile(os.path.join(save_folder, name + suffix))


if __name__ == '__main__':
    import time
    from snr_test import plot_snr_test
    start_time = time.time()
    #------------------------- main -------------------------#
    import yaml
    config_path = r"F:\research-code\2025\H256+LDPC\configs\channel_codec_config.yaml"
    with open(config_path, "r", encoding="utf-8") as f:
        config = yaml.load(f, yaml.FullLoader)
    config = DictToObject(config)
    image_path = r"D:\data\cifar-10-batches-py\temp_test\33\kodim23_dec.png"
    image = Image.open(image_path)
    bits, shape = image2bits(image)
    start, end, step = -5, 26, 2
    qam_order = 16
    snrs = list(range(start, end, step))
    size = "512x768"
    for snr in snrs:
        config.snr = snr
        simulat_communication = SimulateCommunication(config)
        recieve_bits = simulat_communication(bits)
        # param = get_ldpc_code_params(ldpc_design_filename="./1440.720.txt", compute_matrix=True)
        # recieve_bits = ldpc_qam_awgn(bits, param, snr=5, qam_order=16)
        receive_img = bits2image(recieve_bits, shape)
        img_save_path = r"F:\PCL\month\2025\25-08\experiments\H.265\%s\%s_snr%d.png" \
                        % (size, os.path.basename(image_path).strip("_dec.png"), snr)
        receive_img.save(img_save_path)
        #------------------------- main -------------------------#
        print(f"Total Consuming {time.time() - start_time} secs.")

    origin_image_path = r"D:\video-communication-dataset\图像编解码测试\kodak\archive\kodim23.png"
    snr_image_folder = r"F:\PCL\month\2025\25-08\experiments\H.265\%s" % size
    plot_snr_test(origin_image_path, snr_image_folder, start, end, step, qam_order,
                  save_path=os.path.join(snr_image_folder, "test_%s.png" % size))
    print("Successfully!!!")
    pass
